import torch
from torch import nn
from torch.nn import functional as F


class RNNModel(nn.Module):
    """
    定义一个基于RNN的模型类，用于处理序列数据。

    参数:
    - input_size: 输入数据的特征维度（词向量长度）
    - num_hiddens: 隐藏层单元的数量
    - num_layers: RNN层的数量
    """

    def __init__(self, input_size, num_hiddens, num_layers, **kwargs):
        super(RNNModel, self).__init__(**kwargs)
        # 初始化RNN层
        self.rnn = nn.RNN(input_size, num_hiddens, num_layers)
        # 初始化全连接层，用于输出预测结果
        self.dense = nn.Linear(num_hiddens, input_size)
        # 保存隐藏单元和层数量，供后续使用
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers

    def forward(self, inputs, state):
        """
        前向传播函数，处理输入数据并生成输出。

        参数:
        - inputs: 输入数据序列
        - state: RNN的初始隐藏状态

        返回:
        - output: RNN的输出
        - state: RNN的最终隐藏状态
        """
        # 调整输入数据的维度顺序，以符合RNN输入要求[时间步长, 批量大小, feature]
        inputs = inputs.permute(1, 0)
        # 将输入数据进行独热编码，并转换为浮点数类型

        inputs = F.one_hot(inputs, num_classes=self.dense.out_features).float()
        # 通过RNN层进行数据处理
        output, state = self.rnn(inputs, state)
        # 调整输出数据的形状，以便进行全连接层处理
        output = output.view(-1, self.num_hiddens)
        # 通过全连接层生成最终输出
        output = self.dense(output)
        return output, state
